library(SparseM)
library(spatstat)

#centroid location and region area size
centroid <- as.vector(outer(seq(from = 0.5, by = 1, length = 5), seq(from = 0.5, by = 1, length = 5), 
	function(x, y) complex(real = x, imaginary = y)))
regionArea <- rep(5, 10)

#quadrature points and weights
quadPts <- c(-0.906179845938664,
	-0.538469310105683,
    0.000000000000000,
    0.538469310105683,
    0.906179845938664)
quadWgts <- c(0.236926885056189,
    0.478628670499366,
	0.568888888888889,
	0.478628670499366,
	0.236926885056189)
evalpt <- as.vector(outer(as.vector(outer(quadPts/2, quadPts/2, function(x, y) complex(real = x, imaginary = y))), centroid, "+"))
quadWgts.mat <- t(diag(0.25, length(centroid)) %x% as.vector(outer(quadWgts, quadWgts, "*")))

#true relative risk surface
trueInt <- function(x, y)
{
	100 * (dgamma(x, shape = 1.5, scale = 0.5)/dgamma(0.25, shape = 1.5, scale = 0.5)) *
		(dgamma(y, shape = 1.5, scale = 0.5)/dgamma(0.25, shape = 1.5, scale = 0.5))
}
lambda.true <- trueInt(Re(evalpt), Im(evalpt))/trueInt(0.25,0.25)

#incident matrix
incident.mat <- matrix(as.integer(0), nrow = 10, ncol = 25)
incident.mat[1,1:5] <- 1
incident.mat[2,6:10] <- 1
incident.mat[3,11:15] <- 1
incident.mat[4,16:20] <- 1
incident.mat[5,21:25] <- 1
incident.mat[6,seq(1, 21, by = 5)] <- 1
incident.mat[7,seq(2, 22, by = 5)] <- 1
incident.mat[8,seq(3, 23, by = 5)] <- 1
incident.mat[9,seq(4, 24, by = 5)] <- 1
incident.mat[10,seq(5, 25, by = 5)] <- 1

#simulate the poisson point process data
nDataset <- 500
population1 <- function(x, y)
{
	tmp <- floor(y)
	output <- rep(0, length(tmp))
	output[which(tmp == 0)] <- 18
	output[which(tmp == 1)] <- 28
	output[which(tmp == 2)] <- 38
	output[which(tmp == 3)] <- 28
	output[which(tmp == 4)] <- 18
	output
}
population2 <- function(x, y)
{
	tmp <- floor(x)
	output <- rep(0, length(tmp))
	output[which(tmp == 0)] <- 18
	output[which(tmp == 1)] <- 28
	output[which(tmp == 2)] <- 38
	output[which(tmp == 3)] <- 28
	output[which(tmp == 4)] <- 18
	output
}

Z1 <- as.im(population1, owin(xrange = c(0,5), yrange = c(0,5)))
Z2 <- as.im(population2, owin(xrange = c(0,5), yrange = c(0,5)))

sim.data.lst <- list()
sim.exact.data.lst <- list()
for (i in 1:nDataset)
{
	#exact locations
	pp1 <- rpoispp(Z1)
	pp2 <- rpoispp(Z2)
	map1 <- data.frame(map = rep(1, pp1$n), region = floor(pp1$y) %% 5 + 1, x = pp1$x, y = pp1$y)
	map2 <- data.frame(map = rep(2, pp2$n), region = floor(pp2$x) %% 5 + 1, x = pp2$x, y = pp2$y)
	map1 <- map1[order(map1$region),]
	map2 <- map2[order(map2$region),]
	offsets1 <- data.frame(map = 1, region = 1:5, popSize = tapply(map1$region, map1$region, length))
	offsets2 <- data.frame(map = 2, region = 1:5, popSize = tapply(map2$region, map2$region, length))
	map1 <- merge(map1, offsets1)
	map2 <- merge(map2, offsets2)
	tmp1 <- runif(nrow(map1)) <= trueInt(map1$x,map1$y)/trueInt(0.25,0.25)
	tmp2 <- runif(nrow(map2)) <= trueInt(map2$x,map2$y)/trueInt(0.25,0.25)
	map1 <- data.frame(map1, case = tmp1)
	map2 <- data.frame(map2, case = tmp2)
	sim.exact.data.lst[[i]] <- rbind(map1, map2)

	#area censored locations
	map1 <- merge(data.frame(map = 1, region = 1:5, observed = tapply(tmp1, map1$region, sum)), offsets1)
	map2 <- merge(data.frame(map = 2, region = 1:5, observed = tapply(tmp2, map2$region, sum)), offsets2)
	sim.data.lst[[i]] <- rbind(map1, map2)
}
rm("pp1", "pp2", "map1", "map2", "offsets1", "offsets2", "tmp1", "tmp2")

#Gaussian kernel functions
bwSeq <- seq(0.01, 1, by = 0.01)
gaussian.intKernel <- function(centroid, y, bw, cellwidth)
{
	(pnorm(Re(centroid) + cellwidth/2, mean = Re(y), sd = bw) - pnorm(Re(centroid) - cellwidth/2, mean = Re(y), sd = bw)) *
		(pnorm(Im(centroid) + cellwidth/2, mean = Im(y), sd = bw) - pnorm(Im(centroid) - cellwidth/2, mean = Im(y), sd = bw))
}

gaussian.kernel <- function(centroid, y, bw)
{
	xx <- dnorm(Re(centroid), mean = Re(y), sd = bw)
	yy <- dnorm(Im(centroid), mean = Im(y), sd = bw)
	z <- xx * yy
	z <- replace(z, is.nan(z), 0)
	z <- replace(z, z == Inf, 1)
	z
}

intKernel.mat.lst <- list()
S.mat.lst <- list()
for(j in 1:length(bwSeq))
{
	S.mat.lst[[j]]	<- outer(centroid, evalpt, gaussian.kernel, bw = bwSeq[j])
	intKernel.mat.lst[[j]]	<- outer(centroid, evalpt, gaussian.intKernel, bw = bwSeq[j], cellwidth = 1)
}

#expected case count when there is no spatial varying risk surface
offsets <- t(sapply(sim.data.lst, function(x) x$popSize)) %*% diag(1/regionArea) %*% incident.mat

#exact risk estimator
mise.exact.offsets.mat <- matrix(0, nrow = nDataset, ncol = length(bwSeq))
lambda.exact.offsets <- array(0, dim = c(nDataset, length(bwSeq), length(evalpt)))
for(i in 1:nDataset)
{
	exact.data <- sim.exact.data.lst[[i]]
	location <- complex(real = exact.data$x, imaginary = exact.data$y)
	case <- location[exact.data$case]

	for(j in (1:length(bwSeq)))
	{
		S2.mat <- intKernel.mat.lst[[j]]
		denominator <- diag(1/as.vector(offsets[i,] %*% S2.mat))
		lambda <- rep(1, length(case)) %*% outer(case, evalpt, gaussian.kernel, bw = bwSeq[j])
		lambda.exact.offsets[i,j,] <- as.vector(lambda %*% denominator)
		mise.exact.offsets.mat[i,j] <- mean(quadWgts.mat %*% (lambda.exact.offsets[i,j,] - lambda.true)^2)
	}
	cat(i, "\n")
}

save.image("sim_spatial.RData")

#local EM with constant term
tol <- 1e-8
Lambda.ems <- array(0, dim = c(nDataset, length(bwSeq), length(centroid)))
lambda.ems <- array(0, dim = c(nDataset, length(bwSeq), length(evalpt)))
mise.ems.mat <- matrix(0, nrow = nDataset, ncol = length(bwSeq))
for(i in 1:nDataset) 
{
	observed <- sim.data.lst[[i]]$observed
	for(j in 1:length(bwSeq)) 
	{
		K_x.mat <- diag(offsets[i,]) %*% as.matrix.csr(intKernel.mat.lst[[j]] %*% diag(1/as.vector(offsets[i,] %*% intKernel.mat.lst[[j]])))
		K.mat <- K_x.mat %*% t(quadWgts.mat)
		Lambda.old <- as.matrix.csr(matrix(1, nrow = 1, ncol = length(centroid)))
		
		repeat 
		{
			#EM step
			V.inv <- diag(1/as.vector(as.matrix(Lambda.old) %*% t(incident.mat)))
			tmpEM <- (observed %*% V.inv %*% incident.mat) * as.vector(as.matrix(Lambda.old))/offsets[i,]

			#S step
			Lambda <- tmpEM %*% K.mat
			
			if(sum((as.matrix(Lambda) - as.matrix(Lambda.old))^2) > tol) 
			{
				Lambda.old <- Lambda
			}
			else 
			{
				lambda <- tmpEM %*% K_x.mat
				break;
			}
		}

		Lambda.ems[i,j,] <- as.vector(as.matrix(Lambda))
		lambda.ems[i,j,] <- as.vector(as.matrix(lambda))
		mise.ems.mat[i,j] <- mean(quadWgts.mat %*% (lambda.ems[i,j,] - lambda.true)^2)
	}
	cat(i, "\n")
}

save.image("sim_spatial.RData")

#Smoothed EM at the centroids
mise.mid.mat <- matrix(0, nrow = nDataset, ncol = length(bwSeq))
lambda.mid <- array(0, dim = c(nDataset, length(bwSeq), length(evalpt)))
for(i in 1:nDataset) 
{
	Lambda <- Lambda.ems[i,1,]

	for(j in 1:length(bwSeq)) 
	{
		S1.mat <- S.mat.lst[[j]]
		lambda.mid[i,j,] <- as.vector(Lambda %*% S1.mat)
		mise.mid.mat[i,j] <- mean(quadWgts.mat %*% (lambda.mid[i,j,] - lambda.true)^2)
	}
	cat(i, "\n")
}

#compute average MISE
avg.mise.exact.offsets <- apply(mise.exact.offsets.mat, 2, mean)
avg.mise.ems <- apply(mise.ems.mat, 2, mean)
avg.mise.mid <- apply(mise.mid.mat, 2, mean)

save.image("sim_spatial.RData")
